//! ## DBT Cache Table Data Structures in CS2 System
//!
//! This module defines the data structures for the cache table in the CS2 system.
//! The cache table is used to store the translated code and metadata for the guest
//! code. And it is shared between the client and server sides of the CS2 system,
//! supported by [`SharedMemMap`] and its related types/traits.
//!
//! See [`SharedMemMap`] for more details about the shared memory map.

use rustc_hash::FxHasher;

use crate::{
    error::Result,
    sign::CodeSign,
    utils::map::{Borrowed, Config, Key, Owned, SharedMemMap, Val},
};

/// Configuration for the cache table.
pub struct CacheTableConfig;

impl Config<2> for CacheTableConfig {
    const MAGIC: u32 = 0x1AB1EC04;

    const PREFIX: &str = "cs2s-table-";

    type Key = CacheTableKey;

    type Val<'a> = CacheTableVal<'a>;

    type Hasher = FxHasher;
}

/// Key for the cache table.
#[repr(C)]
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct CacheTableKey {
    guest_addr: usize,
    guest_size: usize,
    guest_sign: CodeSign,
}

impl Key for CacheTableKey {
    type Ref<'a> = CacheTableKeyRef<'a>;

    fn as_ref(&self) -> CacheTableKeyRef {
        CacheTableKeyRef {
            guest_addr: self.guest_addr,
            guest_size: self.guest_size,
            guest_sign: &self.guest_sign,
        }
    }

    fn from_ref(key: &Self::Ref<'_>) -> Self {
        Self {
            guest_addr: key.guest_addr,
            guest_size: key.guest_size,
            guest_sign: key.guest_sign.clone(),
        }
    }
}

/// Reference of the cache table key.
#[repr(C)]
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct CacheTableKeyRef<'a> {
    pub guest_addr: usize,
    pub guest_size: usize,
    pub guest_sign: &'a CodeSign,
}

/// Value for the cache table.
#[repr(C)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CacheTableVal<'a> {
    pub host_meta: &'a [u8],
    pub host_code: &'a [u8],
}

impl<'a> Val<'a, 2> for CacheTableVal<'a> {
    type Metadata = ();

    type Content = u8;

    fn metadata(&self) -> &Self::Metadata {
        &()
    }

    fn content(&self) -> [&[Self::Content]; 2] {
        [self.host_meta, self.host_code]
    }

    fn parse(_: &'a Self::Metadata, [host_meta, host_code]: &[&'a [Self::Content]; 2]) -> Self {
        Self {
            host_meta,
            host_code,
        }
    }
}

/// Handle of the cache table for the client.
pub struct CacheTableClientHandle(SharedMemMap<2, CacheTableConfig, Borrowed>);

impl CacheTableClientHandle {
    /// Open the cache table.
    pub fn new(name: &str) -> Result<Self> {
        SharedMemMap::open(name).map(Self)
    }

    /// Check if the cache table needs enlargement.
    #[inline]
    pub fn need_enlargement(&self) -> bool {
        self.0.need_enlargement()
    }

    /// Get the usage of the cache table.
    #[inline]
    pub fn usage(&self) -> usize {
        self.0.usage()
    }

    /// Get the capacity of the cache table.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.0.capacity()
    }

    /// Iterate over the cache table data.
    #[inline]
    pub fn iter_data(&self) -> impl Iterator<Item = (CacheTableKeyRef, CacheTableVal)> {
        self.0.slots().flatten()
    }

    /// Iterate over the cache table slots.
    #[inline]
    pub fn iter_slots(
        &self,
    ) -> impl Iterator<Item = impl Iterator<Item = (CacheTableKeyRef, CacheTableVal)>> {
        self.0.slots()
    }

    /// Lookup the cache table for the given guest code.
    pub fn lookup<'a>(&'a self, key: &'a CacheTableKeyRef) -> Result<Option<CacheTableVal<'a>>> {
        self.0.lookup(key)
    }

    /// Enlarge the cache table.
    pub fn enlarge(&mut self) -> Result<()> {
        self.0.enlarge()
    }
}

/// Handle of the cache table for the server.
pub struct CacheTableServerHandle(SharedMemMap<2, CacheTableConfig, Owned>);

impl CacheTableServerHandle {
    /// Create a new cache table.
    pub fn new(name: &str, nslots: usize, npages: usize) -> Result<Self> {
        SharedMemMap::new(name, nslots, npages).map(Self)
    }

    /// Get the usage of the cache table.
    #[inline]
    pub fn usage(&self) -> usize {
        self.0.usage()
    }

    /// Get the capacity of the cache table.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.0.capacity()
    }

    /// Insert the cache table data.
    pub fn insert(&mut self, key: &CacheTableKeyRef, val: &CacheTableVal) -> Result<()> {
        self.0.insert(key, val)
    }

    /// Enlarge the cache table.
    pub fn enlarge(&mut self, npages: usize) -> Result<()> {
        self.0.enlarge(npages)
    }
}

#[cfg(test)]
mod test {
    use coverage_helper::test;
    use cs2s_testdata::*;

    use crate::error::Error;

    use super::*;

    /// Convenient macro for looking up the cache table from the server side.
    macro_rules! srv_lookup {
        ($handle:expr, $guest_addr:expr, $guest_code:expr) => {
            $handle.0.lookup(&CacheTableKeyRef {
                guest_addr: $guest_addr,
                guest_size: $guest_code.len(),
                guest_sign: &$guest_code.into(),
            })
        };
    }

    /// Convenient macro for looking up the cache table from the client side.
    macro_rules! cli_lookup {
        ($handle:expr, $guest_addr:expr, $guest_code:expr) => {
            $handle.lookup(&CacheTableKeyRef {
                guest_addr: $guest_addr,
                guest_size: $guest_code.len(),
                guest_sign: &$guest_code.into(),
            })
        };
    }

    /// Convenient macro for asserting the cache table lookup results.
    macro_rules! cli_lookup_assert {
        ($handle:expr, [$(
            $guest_addr:expr, $guest_code:expr, $host_meta:expr, $host_code:expr;
        )+]) => {
            $(
                assert_eq!(
                    cli_lookup!($handle, $guest_addr, $guest_code).unwrap(),
                    Some(CacheTableVal {
                        host_meta: &$host_meta,
                        host_code: &$host_code,
                    })
                );
            )+
        };
    }

    /// Convenient macro for asserting the cache table lookup results.
    macro_rules! srv_lookup_assert {
        ($handle:expr, [$(
            $guest_addr:expr, $guest_code:expr, $host_meta:expr, $host_code:expr;
        )+]) => {
            $(
                assert_eq!(
                    srv_lookup!($handle, $guest_addr, $guest_code),
                    Some(CacheTableVal {
                        host_meta: &$host_meta,
                        host_code: &$host_code,
                    })
                );
            )+
        };
    }

    /// Convenient macro for inserting the cache table with assertions.
    macro_rules! srv_insert_assert {
        ($handle:expr, [$(
            $guest_addr:expr, $guest_code:expr, $host_meta:expr, $host_code:expr;
        )+]) => {
            $(
                assert!($handle
                    .insert(&CacheTableKeyRef {
                            guest_addr: $guest_addr,
                            guest_size: $guest_code.len(),
                            guest_sign: &$guest_code.into(),
                        },
                        &CacheTableVal {
                            host_meta: &$host_meta,
                            host_code: &$host_code,
                        },
                    )
                    .is_ok());
                srv_lookup_assert!($handle, [
                    $guest_addr, $guest_code, $host_meta, $host_code;
                ]);
            )+
        };
    }

    #[test]
    fn basic_ops() {
        const NAME: &str = "test-basic-ops";

        // Create a new cache table.
        let mut table = CacheTableServerHandle::new(NAME, 2, 1).unwrap();

        // Should not find any meta/code in an empty table.
        assert!(srv_lookup!(table, 0x400000, GUEST_CODE_0).is_none());

        // Insert three cache data and check if they can be found.
        srv_insert_assert!(table, [
            0x400000, GUEST_CODE_0, HOST_META_0, HOST_CODE_0;
            0x400040, GUEST_CODE_1, HOST_META_1, HOST_CODE_1;
            0x400080, GUEST_CODE_2, HOST_META_2, HOST_CODE_2;
        ]);

        // Confirm that all the inserted data can be found.
        srv_lookup_assert!(table, [
            0x400000, GUEST_CODE_0, HOST_META_0, HOST_CODE_0;
            0x400040, GUEST_CODE_1, HOST_META_1, HOST_CODE_1;
            0x400080, GUEST_CODE_2, HOST_META_2, HOST_CODE_2;
        ]);

        // Insert more data to test the collision handling.
        srv_insert_assert!(table, [
            0x4000C0, GUEST_CODE_3, HOST_META_3, HOST_CODE_3;
        ]);

        // Check if all the data can still be found.
        srv_lookup_assert!(table, [
            0x400000, GUEST_CODE_0, HOST_META_0, HOST_CODE_0;
            0x400040, GUEST_CODE_1, HOST_META_1, HOST_CODE_1;
            0x400080, GUEST_CODE_2, HOST_META_2, HOST_CODE_2;
            0x4000C0, GUEST_CODE_3, HOST_META_3, HOST_CODE_3;
        ]);
    }

    #[test]
    fn space_enlargement() {
        const NAME: &str = "test-space-enlargement";

        // Create a new cache table.
        let mut table = CacheTableServerHandle::new(NAME, 256, 1).unwrap();

        // Open the cache table from the client side.
        let mut ro_table = CacheTableClientHandle::new(NAME).unwrap();

        let mut enlargement_count = 0;

        for i in 0.. {
            let guest_addr = 0x400000 + i * 0x40;
            let guest_code = GUEST_CODES[i % GUEST_CODES.len()];
            let host_meta = HOST_METAS[i % HOST_METAS.len()];
            let host_code = HOST_CODES[i % HOST_CODES.len()];
            if let Err(Error::MemNotEnough { .. }) = table.insert(
                &CacheTableKeyRef {
                    guest_addr,
                    guest_size: guest_code.len(),
                    guest_sign: &guest_code.into(),
                },
                &CacheTableVal {
                    host_meta,
                    host_code,
                },
            ) {
                enlargement_count += 1;
                table.enlarge(1).unwrap();
                srv_insert_assert!(table, [
                    guest_addr, guest_code, host_meta, host_code;
                ]);
                assert!(ro_table.need_enlargement());
                ro_table.enlarge().unwrap();
            }
            assert!(ro_table.usage() <= ro_table.capacity());
            cli_lookup_assert!(ro_table, [
                guest_addr, guest_code, host_meta, host_code;
            ]);

            if enlargement_count >= 100 {
                break;
            }
        }
    }

    #[test]
    fn insert_existing_cache() {
        const NAME: &str = "test-insert-existing-cache";

        // Create a new cache table.
        let mut table = CacheTableServerHandle::new(NAME, 1, 1).unwrap();

        // Insert a new cache data.
        let usage = table.usage();
        srv_insert_assert!(table, [
            0x400000, GUEST_CODE_0, HOST_META_0, HOST_CODE_0;
        ]);
        assert!(table.usage() > usage);

        // Insert the same cache data again.
        let usage = table.usage();
        srv_insert_assert!(table, [
            0x400000, GUEST_CODE_0, HOST_META_0, HOST_CODE_0;
        ]);
        assert_eq!(table.usage(), usage);
    }

    #[test]
    fn table_iter() {
        const NAME: &str = "test-table-iter";

        // Create a new cache table.
        let mut table = CacheTableServerHandle::new(NAME, 2, 16).unwrap();

        // Open the cache table from the client side.
        let ro_table = CacheTableClientHandle::new(NAME).unwrap();

        // Insert some data to the server side.
        for i in 0..128 {
            let guest_addr = 0x400000 + i * 0x40;
            let guest_code = GUEST_CODES[i % GUEST_CODES.len()];
            let host_meta = HOST_METAS[i % HOST_METAS.len()];
            let host_code = HOST_CODES[i % HOST_CODES.len()];
            table
                .insert(
                    &CacheTableKeyRef {
                        guest_addr,
                        guest_size: guest_code.len(),
                        guest_sign: &guest_code.into(),
                    },
                    &CacheTableVal {
                        host_meta,
                        host_code,
                    },
                )
                .unwrap();
        }

        // Iterate over the cache data and check if they can be found.
        for i in 0..128 {
            let guest_addr = 0x400000 + i * 0x40;
            let guest_code = GUEST_CODES[i % GUEST_CODES.len()];
            let host_meta = HOST_METAS[i % HOST_METAS.len()];
            let host_code = HOST_CODES[i % HOST_CODES.len()];
            let existence = ro_table.iter_data().any(|(key, val)| {
                key.guest_addr == guest_addr
                    && key.guest_size == guest_code.len()
                    && key.guest_sign == &guest_code.into()
                    && val.host_meta == host_meta
                    && val.host_code == host_code
            });
            assert!(existence);

            let existence = ro_table.iter_slots().any(|mut slot| {
                slot.any(|(key, val)| {
                    key.guest_addr == guest_addr
                        && key.guest_size == guest_code.len()
                        && key.guest_sign == &guest_code.into()
                        && val.host_meta == host_meta
                        && val.host_code == host_code
                })
            });
            assert!(existence);
        }

        // Assert that the iterator can be exhausted.
        assert!(!ro_table.iter_data().any(|_| false));
        assert!(!ro_table.iter_slots().any(|_| false));
        assert!(!ro_table.iter_slots().any(|mut slot| slot.any(|_| false)));
    }
}
